# !/usr/bin/env python
# -*- coding: utf-8 -*-


"""
ip地址白名单黑名单

加入redis中

每次请求需要查看IP地址是否在白名单中（这需要请求中间件完成）

IP白名单保存在redis中，类型为set
"""

from starlette.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from app.db.redis_utils import red
from app.core.log import logger
from app.core.log_request import logger as logger_api


class RequestLogMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        logger_api.info(request.url)
        response = await call_next(request)
        return response


class IPFirewallMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        ip_or_host = request.client.host
        _header = request.headers

        if 'X-Forwarded-For' in _header:
            ip_or_host = _header['X-Forwarded-For']
        # 记录请求日志
        logger_api.info(ip_or_host)

        if await red.execute('sismember', 'open_data_api_ip_firewall', ip_or_host) == 0:
            logger.error("ip:{ip}  IP认证失败，访问被拒绝，".format(ip=ip_or_host))
            return JSONResponse("Your IP is not allowed to access", status_code=401)

        response = await call_next(request)
        response.headers['Custom'] = 'Example'
        return response
